"""
ConnLLM安全工具函数
包含日志安全和敏感信息处理等功能
"""
import re
import logging
from typing import Dict, Any, List, Union, Callable, Set, Pattern

# 需要屏蔽的字段名称集合
SENSITIVE_FIELDS = {
	"api_key", "apikey", "secret", "password", "token", "authorization", 
	"auth", "private_key", "client_secret", "access_token", "refresh_token"
}

# 需要屏蔽的头部前缀
SENSITIVE_HEADERS = {
	"x-api-key", "authorization", "api-key", "auth"
}


def mask_sensitive_text(text: str, mask_char: str = "*") -> str:
	"""
	屏蔽文本中的敏感信息
	
	Args:
		text: 原始文本
		mask_char: 用于屏蔽的字符
		
	Returns:
		屏蔽后的文本
	"""
	if not text:
		return text
		
	# 屏蔽API密钥格式
	patterns = [
		# Anthropic密钥格式: sk-ant-api03-...
		(r'sk-ant-[a-zA-Z0-9]{8,}', r'sk-ant-***'),
		# OpenAI密钥格式: sk-...
		(r'sk-[a-zA-Z0-9]{8,}', r'sk-***'),
		# OpenRouter密钥格式: sk-or-...
		(r'sk-or-[a-zA-Z0-9]{8,}', r'sk-or-***'),
		# Moonshot密钥格式: sk-ms-...
		(r'sk-ms-[a-zA-Z0-9]{8,}', r'sk-ms-***'),
		# 通用Bearer令牌
		(r'Bearer\s+[a-zA-Z0-9._-]{8,}', r'Bearer ***'),
		# 通用API密钥格式 (适用于多种提供商)
		(r'(api[-_]?key|apikey|token)\s*[=:]\s*["\'`]?([a-zA-Z0-9._-]{8,})["\'`]?', r'\1=***')
	]
	
	masked = text
	for pattern, replacement in patterns:
		masked = re.sub(pattern, replacement, masked, flags=re.IGNORECASE)
		
	return masked


def mask_sensitive_fields(data: Any, fields_to_mask: Union[Set[str], None] = None, 
						 mask_char: str = "*", path: str = "") -> Any:
	"""
	递归屏蔽字典或列表中的敏感字段
	
	Args:
		data: 要处理的数据
		fields_to_mask: 要屏蔽的字段名集合，如果为None则使用默认集合
		mask_char: 用于屏蔽的字符
		path: 当前处理的字段路径(用于递归)
		
	Returns:
		屏蔽后的数据副本
	"""
	if fields_to_mask is None:
		fields_to_mask = SENSITIVE_FIELDS
		
	# 处理不同类型的数据
	if data is None:
		return None
	elif isinstance(data, dict):
		masked_dict = {}
		for key, value in data.items():
			curr_path = f"{path}.{key}" if path else key
			
			# 判断是否为敏感字段
			if (key.lower() in fields_to_mask or
				any(curr_path.lower().endswith(f".{field}") for field in fields_to_mask) or
				any(h.lower() in key.lower() for h in SENSITIVE_HEADERS)):
				# 敏感字段，根据类型进行屏蔽
				if isinstance(value, str) and value:
					# 字符串，保留前后3个字符
					if len(value) > 8:
						masked_dict[key] = value[:3] + mask_char * (len(value) - 6) + value[-3:]
					else:
						masked_dict[key] = mask_char * len(value)
				elif value is not None:
					# 其他非空值，用屏蔽符号代替
					masked_dict[key] = f"{mask_char * 8}"
				else:
					# None值，保持不变
					masked_dict[key] = None
			else:
				# 非敏感字段，递归处理
				masked_dict[key] = mask_sensitive_fields(value, fields_to_mask, mask_char, curr_path)
				
		return masked_dict
	elif isinstance(data, list):
		# 列表，递归处理每个元素
		return [mask_sensitive_fields(item, fields_to_mask, mask_char, path) for item in data]
	elif isinstance(data, (str, int, float, bool)) or data is None:
		# 基本类型，直接返回
		return data
	else:
		# 其他类型，转为字符串返回
		return str(data)


class SensitiveInfoFilter(logging.Filter):
	"""
	敏感信息过滤器，用于日志记录
	可以移除日志中的API密钥、令牌等敏感信息
	"""
	
	def __init__(self, patterns: Union[List[Pattern], None] = None):
		"""
		初始化过滤器
		
		Args:
			patterns: 自定义正则表达式模式列表
		"""
		super().__init__()
		self.patterns = patterns or []
		
	def filter(self, record: logging.LogRecord) -> bool:
		"""
		过滤日志记录
		
		Args:
			record: 日志记录
			
		Returns:
			是否保留该记录(总是True，但会修改记录内容)
		"""
		if isinstance(record.msg, str):
			# 处理字符串消息
			record.msg = mask_sensitive_text(record.msg)
			
		# 处理格式化参数
		if hasattr(record, 'args') and record.args:
			args = record.args
			if isinstance(args, dict):
				# 字典参数
				record.args = mask_sensitive_fields(args)
			elif isinstance(args, (list, tuple)):
				# 序列参数
				# 将参数转换为列表，这样可以修改其内容
				new_args = list(args)
				for i, arg in enumerate(new_args):
					if isinstance(arg, (dict, list)):
						new_args[i] = mask_sensitive_fields(arg)
					elif isinstance(arg, str):
						new_args[i] = mask_sensitive_text(arg)
				# 转回原来的类型
				record.args = tuple(new_args) if isinstance(args, tuple) else new_args
				
		return True


def setup_secure_logging() -> None:
	"""
	配置安全日志，过滤敏感信息
	
	修改全局日志处理器，添加敏感信息过滤器
	"""
	# 添加过滤器到根日志记录器
	root_logger = logging.getLogger()
	filter_found = False
	
	# 检查是否已经添加了过滤器
	for existing_filter in root_logger.filters:
		if isinstance(existing_filter, SensitiveInfoFilter):
			filter_found = True
			break
			
	if not filter_found:
		# 创建并添加过滤器
		sensitive_filter = SensitiveInfoFilter()
		root_logger.addFilter(sensitive_filter)
		
		# 同时添加到connllm日志记录器
		connllm_logger = logging.getLogger("connllm")
		connllm_logger.addFilter(sensitive_filter)
		
		# 设置默认日志级别
		if not root_logger.handlers:
			handler = logging.StreamHandler()
			formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
			handler.setFormatter(formatter)
			root_logger.addHandler(handler)
			
		# 如果没有设置级别，设置为INFO
		if root_logger.level == logging.NOTSET:
			root_logger.setLevel(logging.INFO)
